package com.example.harmonet.harmtorch;

public class MaxPool2d implements Layer {

    int[] kernel_size;
    int[] stride;
    int[] padding;

    public MaxPool2d(int[] kernel_size, int[] stride, int[] padding) {
        this.kernel_size = kernel_size.clone();
        this.stride = stride.clone();
        this.padding = padding.clone();
    }

    private float getMax(Tensor in, int n, int i, int h, int w) throws Exception {
        float max = -Float.MAX_VALUE;
        for(int h2 = 0; h2 < kernel_size[0]; ++h2) {
            for(int w2 = 0; w2 < kernel_size[1]; ++w2) {
                if(max < in.tensor()[in.pos(n, i, h * stride[0] + h2, w * stride[1] + w2)]) {
                    max = in.tensor()[in.pos(n, i, h * stride[0] + h2, w * stride[1] + w2)];
                }
            }
        }
        return max;
    }

    private void computeWH(Tensor in, Tensor out, int n, int i, int[] out_dim) throws Exception {
        for(int h = 0; h < out_dim[2]; ++h) {
            for(int w = 0; w < out_dim[3]; ++w) {
                out.tensor()[out.pos(n, i, h, w)] = getMax(in, n, i, h, w);
            }
        }
    }

    @Override
    public Tensor forward(Tensor in) throws Exception {
        int[] out_dim = new int[4];
        out_dim[0] = in.dim()[0];
        out_dim[1] = in.dim()[1];
        out_dim[2] = (in.dim()[2] + 2 * padding[0] - kernel_size[0]) / stride[0] + 1;
        out_dim[3] = (in.dim()[3] + 2 * padding[1] - kernel_size[1]) / stride[1] + 1;
        Tensor out = new Tensor(out_dim);
        for(int n = 0; n < out.dim()[0]; ++n) {
            for(int i = 0; i < out_dim[1]; ++i) {
                computeWH(in, out, n, i, out_dim);
            }
        }
        return out;
    }

    @Override
    public int getParam() {
        return 0;
    }

    @Override
    public void init(float[] param) {
    }
}
